package com.prac.mq.config;

import java.util.HashMap;
import java.util.Map;
import java.util.function.Consumer;

import lombok.extern.slf4j.Slf4j;
import org.apache.rocketmq.client.producer.LocalTransactionState;
import org.apache.rocketmq.client.producer.TransactionListener;
import org.apache.rocketmq.client.producer.TransactionMQProducer;
import org.apache.rocketmq.common.message.Message;
import org.apache.rocketmq.common.message.MessageExt;
import org.apache.rocketmq.spring.autoconfigure.RocketMQAutoConfiguration;
import org.apache.rocketmq.spring.core.RocketMQTemplate;
import org.springframework.beans.BeansException;
import org.springframework.boot.autoconfigure.AutoConfigureAfter;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.annotation.Configuration;

@Configuration
@AutoConfigureAfter({RocketMQAutoConfiguration.class})
@Slf4j
public class RocketConfig implements ApplicationContextAware {

    private static Map<Integer, LocalTransactionState> map = new HashMap<>(8);
    private static final int UNKNOW = 0;
    private static final int COMMIT_MESSAGE = 1;
    private static final int ROLLBACK_MESSAGE = 2;
    private static final Map<String, Integer> redis = new HashMap<>();

    @Override
    public void setApplicationContext(ApplicationContext context) throws BeansException {
        RocketMQTemplate mqTemplate = context.getBean(RocketMQTemplate.class);
        /**
         * 获取事务消息生产者，统一添加事务监听器
         */
        TransactionMQProducer producer = (TransactionMQProducer) mqTemplate.getProducer();
        producer.setTransactionListener(new TransactionListener() {
            @Override
            public LocalTransactionState executeLocalTransaction(Message msg, Object arg) {
                /**
                 * rocket消息生成时会自动生成一个唯一id和唯一TransactionId
                 */
                //msg.setTransactionId(!Objects.isNull(msg.getTransactionId()) ? msg.getTransactionId() : IdUtil.randomUUID());
                int state;
                if (arg instanceof Consumer) {
                    try {
                        ((Consumer<String>) arg).accept(new String(msg.getBody()));
                        state = COMMIT_MESSAGE;
                    } catch (Exception e) {
                        log.warn("trans-msgId:{},executeLocalTransaction error", msg.getTransactionId(), e);
                        state = ROLLBACK_MESSAGE;
                    }
                } else {
                    log.warn("executeLocalTransaction error, arg is not Consumer");
                    state = ROLLBACK_MESSAGE;
                }
                /**
                 * redis记录 事务执行状态
                 */
                log.info("redis存储transactionId:{}", msg.getTransactionId());
                redis.put(msg.getTransactionId(), state);
                return map.get(state);
            }

            @Override
            public LocalTransactionState checkLocalTransaction(MessageExt msg) {
                log.info("回查checkLocalTransaction id:{}", msg.getTransactionId());
                Integer state = redis.get(msg.getTransactionId());
                return map.get(state);
            }
        });
        log.info("=====>RocketMQTemplate init TransactionListener success");
    }


    static {
        map.put(UNKNOW, LocalTransactionState.UNKNOW);
        map.put(COMMIT_MESSAGE, LocalTransactionState.COMMIT_MESSAGE);
        map.put(ROLLBACK_MESSAGE, LocalTransactionState.ROLLBACK_MESSAGE);
    }

}
